import threading
from collections import deque
from typing import Any, Callable, Iterable, Optional, Union

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.widgets import Button

from .repeating_timer import RepeatingTimer
from .subprocess_handle import Connection, SubprocessHandle
from .utils import convert_to_iterator

HistPoint = Union[float, Iterable[float]]
HistBufferData = Optional[tuple[np.ndarray, ...]]
HistViewData = Optional[tuple[tuple[np.ndarray, np.ndarray], ...]]


def show_summary(data: HistViewData, title: Optional[str]):
    ys_histogram = data

    plt.figure(figsize=(8, 4))
    axes: plt.Axes = plt.gca()
    axes.grid()

    if title is not None:
        axes.set_title(title)

    # draw bars (and calculate min max)
    max_height = 0
    min_x, max_x = None, None
    for yi, (y_heights, y_edges) in enumerate(ys_histogram):
        axes.bar(y_edges[:-1], y_heights, y_edges[1:] - y_edges[:-1],
            align='edge', edgecolor="black", linewidth=1, zorder=2 + (1 - yi / len(ys_histogram)))
        max_height = max(max_height, np.max(y_heights))
        min_x = min(min_x, min(y_edges)) if min_x is not None else min(y_edges)
        max_x = max(max_x, max(y_edges)) if max_x is not None else max(y_edges)

    # adjust xy range
    padding = (max_x - min_x) / 32
    axes.set_xlim(min_x - padding, max_x + padding)
    axes.set_ylim(0, max_height * 1.1)

    plt.show()


class AccumulateHistBuffer:
    def __init__(self, size: Optional[int] = 1024):
        self._lock = threading.RLock()
        self._condition = threading.Condition(self._lock)

        self._version = 0
        self._ys = deque[np.ndarray](maxlen=size)

    def append(self, ys: HistPoint):
        ys = np.asarray(ys, dtype=np.float32)
        if ys.ndim == 0:
            ys = ys.reshape([1])
        elif ys.ndim == 1:
            ys = ys.reshape([-1])
        else:
            raise ValueError(f"Expected ys.ndim = 0 or 1, got ys.shape={ys.shape}")

        with self._lock:
            self._version += 1
            self._ys.append(ys)
            self._condition.notify_all()

    def clear(self):
        with self._lock:
            self._version += 1
            self._ys.clear()
            self._condition.notify_all()

    def read(self) -> tuple[int, HistBufferData]:
        with self._lock:
            if len(self._ys) == 0:
                return self._version, None
            return self._version, tuple(self._ys)

    def wait_and_read(self, version: int) -> tuple[int, HistBufferData]:
        with self._lock:
            while version == self._version:
                self._condition.wait()
            return self.read()


def convert_to_view_data(data: HistBufferData, bins: int) -> HistViewData:
    if data is None:
        return None
    ys = np.transpose(data)  # [ys_n, seq_n]
    return tuple(np.histogram(y, bins=bins) for y in ys)


class AccumulateHistView:
    def __init__(self,
        figure: plt.Figure, *,
        title: Optional[str] = None,
    ):
        self._figure = self._init_figure(figure)
        self._axes = self._init_axes(figure, title)
        self._play_button, self._clear_button = self._init_buttons(figure)
        self._play_button.on_clicked(lambda _: self._on_play_clicked())
        self._clear_button.on_clicked(lambda _: self._on_clear_clicked())

        self._ys_bars: Optional[list[Iterable[plt.Rectangle]]] = None

        self._paused = True
        self.request_pause: Optional[Callable[[], Any]] = None
        self.request_resume: Optional[Callable[[], Any]] = None
        self.request_clear: Optional[Callable[[], Any]] = None

    @staticmethod
    def _init_figure(figure: plt.Figure):
        figure.set_size_inches(8, 4)
        return figure

    @staticmethod
    def _init_axes(figure: plt.Figure, title: Optional[str] = None) -> plt.Axes:
        axes = figure.gca()
        axes.grid()

        if title is not None:
            axes.set_title(title)

        return axes

    @staticmethod
    def _init_buttons(figure: plt.Figure) -> tuple[plt.Button, plt.Button]:
        figure.subplots_adjust(0.1, 0.1, 0.95, 0.85)
        play_button = Button(figure.add_axes([0.01, 0.9, 0.09, 0.09]), "resume")
        clear_button = Button(figure.add_axes([0.11, 0.9, 0.09, 0.09]), "clear")
        return play_button, clear_button

    @staticmethod
    def _init_bars(axes: plt.Axes, ys_n: int, bars_n: int) -> list[tuple[plt.Rectangle, ...]]:
        return [axes.bar(
            np.zeros([bars_n]), np.zeros([bars_n]),
            edgecolor="black", linewidth=1,
            zorder=2 + (ys_n - yi) / ys_n
        ) for yi in range(ys_n)]

    def update(self, data: HistViewData):
        if data is None:  # clear
            self._axes.cla()
            self._axes.grid()
            self._ys_bars = None
            self._figure.canvas.draw_idle()
            return

        ys_histogram = data
        if self._ys_bars is None:
            ys_n = len(ys_histogram)
            if ys_n == 0:
                return
            bars_n = len(ys_histogram[0][0])
            if bars_n == 0:
                return
            self._ys_bars = self._init_bars(self._axes, ys_n, bars_n)

        max_height = 0
        min_x, max_x = None, None
        for y_bars, (y_heights, y_edges) in zip(self._ys_bars, ys_histogram):
            y_widths = y_edges[1:] - y_edges[:-1]
            for bar, x, width, height in zip(y_bars, y_edges[:-1], y_widths, y_heights):
                bar.set_x(x)
                bar.set_width(width)
                bar.set_height(height)
            max_height = max(max_height, np.max(y_heights))
            min_x = min(min_x, min(y_edges)) if min_x is not None else min(y_edges)
            max_x = max(max_x, max(y_edges)) if max_x is not None else max(y_edges)

        # adjust y range
        self._axes.set_ylim(0, max_height * 1.1)

        # adjust both x range
        padding = (max_x - min_x) / 32
        self._axes.set_xlim(min_x - padding, max_x + padding)

        # redraw
        self._figure.canvas.draw_idle()

    def pause(self):
        if self._paused:
            return
        self._paused = True
        self._play_button.label.set_text("resume")

    def resume(self):
        if not self._paused:
            return
        self._paused = False
        self._play_button.label.set_text("pause")

    def _on_play_clicked(self):
        if self._paused:
            if self.request_resume:
                self.request_resume()
        else:
            if self.request_pause:
                self.request_pause()

    def _on_clear_clicked(self):
        if self.request_clear:
            self.request_clear()


def subprocess_func(
    connection: Connection,
    frame_interval: float,
    bins: int,
    title: Optional[str],
    buffer_size: Optional[int]
):
    buffer = AccumulateHistBuffer(buffer_size)

    matplotlib.use('QtAgg')
    figure = plt.figure()
    view = AccumulateHistView(figure, title=title)
    view.request_pause = lambda: connection.send(('ui', "pause"))
    view.request_resume = lambda: connection.send(('ui', "resume"))
    view.request_clear = lambda: connection.send(('ui', "clear"))

    # data timer

    def tick_func():
        nonlocal view_version
        view_version, data = buffer.wait_and_read(view_version)
        view_data = convert_to_view_data(data, bins)
        view.update(view_data)

    view_version = -1
    timer = RepeatingTimer(frame_interval, tick_func)
    timer.start()

    # recv thread

    def on_recv_func(message: tuple[str, Any]):
        typ, data = message
        if typ == 'data':
            buffer.append(data)
        if typ == 'ui':
            if data == "pause":
                view.pause()
            elif data == "resume":
                view.resume()
            elif data == "clear":
                buffer.clear()

    def thread_func():
        while True:
            try:
                message = connection.recv()
            except EOFError:
                break
            on_recv_func(message)

    thread = threading.Thread(target=thread_func, daemon=True)
    thread.start()

    # window
    plt.show()


class AccumulateHistPopup(SubprocessHandle):
    def __init__(self,
        producer: Union[Iterable[HistPoint], Callable[[], HistPoint]], *,
        produce_interval: float = 1 / 128,
        frame_interval: float = 1 / 32,
        bins: int = 20,
        title: Optional[str] = None,
        start: bool = True,
        auto_close: bool = False,
        buffer_size: Optional[int] = 1024,
    ):
        """ Show histograms of accumulated data from `producer` in a popup window

        :param producer: an iterator (generator) or callable that produces data point with shape (ys_n,)
        :param produce_interval: interval between calls of `producer`
        :param frame_interval: interval between gui updates
        :param bins: number of bins in the histogram
        :param title: the title of the plot
        :param start: whether to start immediately
        :param auto_close: close the popup window when producer is exhausted
        :param buffer_size: the number of buffered data points. if None, the buffer is unbounded. 
        """
        super().__init__(
            subprocess_func,
            frame_interval=frame_interval,
            bins=bins, title=title,
            buffer_size=buffer_size)

        # configs for summary
        self._bins = bins
        self._title = title

        # data source and buffer
        self._iterator = convert_to_iterator(producer)
        self._buffer = AccumulateHistBuffer(buffer_size)

        # data timer
        def timer_tick_func():
            try:
                point = next(self._iterator)
            except StopIteration:
                self._timer.stop()
                if auto_close:
                    self.close()
                return
            try:
                self._send(('data', point))
            except (OSError, BrokenPipeError):
                self._timer.stop()
                return
            self._buffer.append(point)

        self._timer = RepeatingTimer(produce_interval, timer_tick_func)

        # recv thread
        def on_recv_func(message: tuple[str, Any]):
            typ, data = message
            if typ == 'ui':
                if data == "pause":
                    self.pause()
                if data == "resume":
                    self.resume()
                if data == "clear":
                    self.clear()

        def thread_func():
            while True:
                try:
                    message = self._recv()
                except EOFError:
                    break
                on_recv_func(message)

        self._thread = threading.Thread(target=thread_func)
        self._thread.start()

        # start
        if start:
            self.resume()

    def resume(self):
        self._timer.start()
        self._send(('ui', "resume"))

    def pause(self):
        self._timer.stop()
        self._send(('ui', "pause"))

    def clear(self):
        self._buffer.clear()
        self._send(('ui', "clear"))

    def close(self):
        super().close()
        self._timer.stop()

    def wait(self):
        super().wait()
        self._timer.join()
        self._thread.join()

    def show_summary(self):
        _, data = self._buffer.read()
        view_data = convert_to_view_data(data, self._bins)
        show_summary(view_data, title=self._title)

    def wait_and_show_summary(self):
        self.wait()
        self.show_summary()

    def close_and_wait_and_show_summary(self):
        self.close()
        self.wait()
        self.show_summary()
